(* ===== 类型定义 ===== *)
(* 抽象语法树类型 *)
type ast =
  | Atom of string    (* 原子类型：符号或数字 *)
  | List of ast list  (* 列表类型：包含多个AST节点 *)

(* 值类型：求值结果 *)
type value =
  | Symbol of string  (* 符号值 *)
  | Number of int     (* 整数值 *)
  | Pair of value * value (* 点对：car和cdr *)
  | Nil               (* 空列表 *)
  | Bool of bool      (* 布尔值 *)
  | Lambda of string list * ast * env
  | Label of string * value

(* 环境类型：存储变量绑定 *)
and env = (string * value) list

(* ===== 词法分析器 ===== *)
type token = LParen | RParen | AtomToken of string | QuoteToken

(* 辅助函数：判断字符是否为空白 *)
let is_space c = 
  c = ' ' || c = '\t' || c = '\n' || c = '\r'

let rec drop_while p = function
| [] -> []
| x :: xs when p x -> drop_while p xs
| xs -> xs

let break_on p lst =
  let rec loop acc = function
    | [] -> (List.rev acc, [])
    | x :: xs when not (p x) -> loop (x :: acc) xs
    | xs -> (List.rev acc, xs)
  in
  loop [] lst

let rec tokenize' str =
  match str with
  | [] -> []
  | '(' :: cs -> 
      LParen :: tokenize' (drop_while is_space cs)
  | ')' :: cs -> 
      RParen :: tokenize' (drop_while is_space cs)
  | '\'' :: cs -> 
      QuoteToken :: tokenize' (drop_while is_space cs)
  | cs ->
      let atom, rest = 
        break_on (fun c -> 
          (is_space c || c = '(' || c = ')' || c = '\'')) cs 
      in
      AtomToken (String.of_seq (List.to_seq atom)) :: 
      tokenize' (drop_while is_space rest)

let tokenize str = 
  tokenize' (List.of_seq (String.to_seq str))

(* ===== 语法分析器 ===== *)
let parse tokens =
  (* 解析列表：直到遇到右括号 *)
  let rec parse_list tokens acc =
    match tokens with
    | [] -> failwith "Unexpected end of input"  (* 错误：输入意外结束 *)
    | RParen :: rest -> (List.rev acc, rest)  (* 结束列表：返回反转的AST节点和剩余token *)
    | _ ->  (* 处理列表元素 *)
        let (ast, rest) = parse_expr tokens in  (* 递归解析表达式 *)
        parse_list rest (ast :: acc)  (* 将解析结果加入累积列表 *)
  
  (* 解析单个表达式 *)
  and parse_expr tokens =
    match tokens with
    | [] -> failwith "Unexpected end of input"  (* 错误：输入意外结束 *)
    | LParen :: rest ->  (* 左括号：开始解析列表 *)
        let (list, rest) = parse_list rest [] in  (* 解析子列表 *)
        (List list, rest)  (* 返回List节点和剩余token *)
    | AtomToken s :: rest -> (Atom s, rest)  (* 原子token：创建Atom节点 *)
    | RParen :: _ -> failwith "Unexpected ')'"  (* 错误：未预期的右括号 *)
    | QuoteToken :: rest ->  (* 引号：处理引用表达式 *)
        let (quoted_ast, rest_after_expr) = parse_expr rest in  (* 解析被引用的表达式 *)
        (List [Atom "quote"; quoted_ast], rest_after_expr)  (* 构建(quote ...)形式 *)
  
  in
  let (ast, rest) = parse_expr tokens in  (* 开始解析表达式 *)
  if rest <> [] then failwith "Unexpected tokens after expression"  (* 检查是否有剩余token *)
  else ast  (* 返回AST *)

(* ===== 辅助函数 ===== *)
(* 尝试将字符串解析为整数，否则视为符号 *)
let try_int s =
  try Number (int_of_string s)  (* 成功：返回Number *)
  with _ -> Symbol s  (* 失败：返回Symbol *)

(* 检查值是否为真（Lisp真值规则） *)
let is_true = function
  | Bool false | Nil -> false  (* false和空列表为假 *)
  | _ -> true  (* 其他值为真 *)

(* 将AST直接转换为值（用于quote特殊形式） *)
let rec ast_to_value = function
  | Atom s -> try_int s  (* 原子：尝试转为数字或符号 *)
  | List lst ->  (* 列表：递归构建Pair链表 *)
      let rec build_list = function
        | [] -> Nil  (* 空列表：返回Nil *)
        | hd :: tl -> Pair (ast_to_value hd, build_list tl)  (* 构建点对 *)
      in
      build_list lst

(* 创建初始环境：包含六大基本函数 *)
let initial_env () : env = []

(* ===== 求值器 ===== *)
let rec eval ast env =
  match ast with
  | Atom s ->  (* 原子求值 *)
      begin match List.assoc_opt s !env with
      | Some v -> v  (* 在环境中找到绑定值 *)
      | None -> try_int s  (* 未找到：尝试转为数字或符号 *)
      end
  
  | List [] -> Nil  (* 空列表求值为Nil *)
  
  | List [Atom "quote"; arg] ->  (* quote特殊形式 *)
      ast_to_value arg  (* 直接转换参数为值 *)

  | List [Atom "atom"; arg] ->
    let eval_atom value = 
      begin match value with
      | Symbol _ | Number _ | Nil -> Bool true
      | _ -> Bool false
      end
    in eval_atom (eval arg env)

  | List [Atom "eq"; lhs; rhs] ->
    let rec eval_eq lhs rhs = 
      begin match (lhs, rhs) with
      | Symbol a, Symbol b -> Bool (a = b)
      | Number a, Number b -> Bool (a = b)
      | Nil, Nil -> Bool true
      | Pair (a1, b1), Pair(a2, b2) ->
        if (eval_eq a1 a2) = Bool true
          then (eval_eq b1 b2)
        else Bool false
      | _ -> Bool false
    end in
    eval_eq (eval lhs env) (eval rhs env)

  | List [Atom "car"; arg] ->
    begin match (eval arg env) with
    | Pair (car, _) -> car
    | Nil -> Nil
    | _ -> failwith "car: expected pair"
    end

  | List [Atom "cdr"; arg] ->
    begin match (eval arg env) with
    | Pair (_, cdr) -> cdr
    | Nil -> Nil
    | _ -> failwith "car: expected pair"
    end

  | List [Atom "cons"; lhs; rhs] ->
    Pair (eval lhs env, eval rhs env)
    
  (* cond特殊形式处理 *)
  | List (Atom "cond" :: clauses) -> 
      let rec eval_clauses = function
        | [] -> failwith "cond: no true clause"  (* 没有真值子句 *)
        | clause :: rest ->
            match clause with
            | List [test; result] ->  (* 标准子句形式 *)
                let test_val = eval test env in  (* 求值测试表达式 *)
                if is_true test_val then eval result env  (* 测试为真：求值结果 *)
                else eval_clauses rest  (* 测试为假：尝试后续子句 *)
            | _ -> failwith "cond: malformed clause"  (* 无效子句格式 *)
      in
      eval_clauses clauses  (* 开始处理子句 *)

  (* lambda *)
  | List [Atom "lambda"; arg; body] ->
    begin match arg with
    | List params ->
      let param_names = List.map (function 
            | Atom s -> s 
            | _ -> failwith "lambda: parameter must be atom"
          ) params in 
      Lambda (param_names, body, !env)
    | _ -> failwith "Lambda with atom param list."
    end
  
  (* define *)
  | List [Atom "label"; Atom name; lambda] -> 
    Label (name, eval lambda env)

  (* 处理define特殊形式 *)
  | List [Atom "define"; Atom name; expr] ->
    let value = eval expr env in
    env := (name, value) :: !env; (* 更新环境 *)
    value

  | List [Atom "+"; lhs; rhs] ->
    begin match (eval lhs env, eval rhs env) with
      | (Number a, Number b) -> Number (a + b)
      | _ -> failwith "+ not called with two numbers"
    end

  | List [Atom "-"; lhs; rhs] ->
    begin match (eval lhs env, eval rhs env) with
      | (Number a, Number b) -> Number (a - b)
      | _ -> failwith "- not called with two numbers"
    end
  
  | List [Atom "*"; lhs; rhs] ->
    begin match (eval lhs env, eval rhs env) with
      | (Number a, Number b) -> Number (a * b)
      | _ -> failwith "* not called with two numbers"
    end

  | List [Atom "/"; lhs; rhs] ->
    begin match (eval lhs env, eval rhs env) with
      | (Number a, Number b) -> Number (a / b)
      | _ -> failwith "/ not called with two numbers"
    end

  | List (func :: args) ->  (* 函数应用 *)
      let func_val = eval func env in  (* 求值函数表达式 *)
      let arg_vals = List.map (fun arg -> eval arg env) args in  (* 求值所有参数 *)
      begin match func_val with
      | Lambda (params, body, closure_env) -> 
        let new_env = ref (List.combine params arg_vals @ closure_env) in
        eval body new_env
      | Label (name, (Lambda (params, body, closure_env) as lambda)) ->
        let new_env = ref ((name, Label (name, lambda)) :: (List.combine params arg_vals @ closure_env)) in
        eval body new_env
      | _ -> failwith "Cannot apply non-function"  (* 非函数不可应用 *)
      end

(* ===== 打印函数 ===== *)
let rec string_of_value = function
  | Symbol s -> s
  | Number n -> string_of_int n
  | Bool true -> "t"  (* Lisp传统真值表示 *)
  | Bool false -> "()"  (* Lisp传统假值表示 *)
  | Nil -> "()"  (* 空列表表示 *)
  | Pair (car, cdr) ->  (* 点对/列表打印 *)
      let rec print_list = function
        | Pair (hd, tl) ->  (* 标准列表元素 *)
            " " ^ string_of_value hd ^ print_list tl
        | Nil -> ""  (* 列表结束 *)
        | v -> " . " ^ string_of_value v  (* 点对表示 *)
      in
      "(" ^ string_of_value car ^ print_list cdr ^ ")"
  | Lambda (params, _, _) -> "lambda " ^ (String.concat "," params)
  | Label (name, func) -> "label " ^ name 

(* ===== 组合函数 ===== *)
(* 完整流程：字符串 -> token列表 -> AST -> 求值 *)
let s_expr_to_value s env =
  s |> tokenize |> parse |> fun ast -> eval ast env

(* ===== 测试函数 ===== *)
let test_builtins () =
  let env = ref (initial_env ()) in  (* 初始化环境 *)
  
  (* 测试用例列表（表达式, 期望结果） *)
  let tests = [
    "(quote a)", "a";
    "(quote (a b c))", "(a b c)";
    "(atom 'a)", "t";
    "(atom '(a b))", "()";
    "(eq 'a 'a)", "t";
    "(eq 'a 'b)", "()";
    "(car '(a b c))", "a";
    "(cdr '(a b c))", "(b c)";
    "(cons 'a '(b c))", "(a b c)";
    "(cond ((eq 1 2) 'a) ((eq 1 1) 'b))", "b"; (* cond测试 *)
    "()", "()";
    "'(1 2 3)", "(1 2 3)";
    "'(a b c)", "(a b c)";
    "(cons 'a (cons 'b ()))", "(a b)";
    "(cond (() 'a) ('t 'b))", "b"; (* 另一个cond测试 *)
    "(+ 3 5)", "8";
    (* Lambda表达式测试 *)
    "((lambda (x) (cons 'a x)) '(b c))", "(a b c)";
    "((lambda (x y) (cons x (cons y ()))) 'a 'b)", "(a b)";
    "((lambda (x) (cons (car x) (cdr x))) '(a b c))", "(a b c)";
    "((lambda (f x) (f x)) (lambda (y) (cons 'a y)) '(b c))", "(a b c)"; (* 高阶函数 *)
    (* Label测试 *)
    "(define add3 (lambda (x) (+ x 3)))", "lambda x";
    "(add3 5)", "8";
    "(define prod (label prod. (lambda (n) (cond ((eq n 0) 1) ('t (* n (prod. (- n 1))))))))", "label prod.";
    "(prod 4)", "24";
    (* Z组合子测试 *)
    "(define Z (lambda (f) ((lambda (x) (f (lambda (v) ((x x) v)))) (lambda (x) (f (lambda (v) ((x x) v)))))))", "lambda f";
    "(define fib. (lambda (g) (lambda (n) (cond ((eq n 0) 0) ((eq n 1) 1) ('t (+ (g (- n 1)) (g (- n 2))))))))", "lambda g";
    "(define fib (Z fib.))", "lambda n";
    "(fib 6)", "8";
  ] in
  
  (* 遍历并执行所有测试 *)
  List.iter (fun (expr, expected) ->
    try
      let result = s_expr_to_value expr env |> string_of_value in
      (* 打印测试结果 *)
      Printf.printf "测试: %-50s 期望: %-20s 结果: %s\n" 
        expr expected result;
      if result = expected then Printf.printf "  [通过]\n"
      else Printf.printf "  [失败]\n"
    with e ->  (* 异常处理 *)
      Printf.printf "测试: %s 出错: %s\n" expr (Printexc.to_string e)
  ) tests

(* ===== REPL 功能 ===== *)
let repl () =
  let env = ref (initial_env ()) in  (* 创建初始环境引用 *)
  print_endline "Lisp REPL (输入 exit 退出)";
  
  let rec loop () =
    print_string "> ";
    flush stdout;
    try
      let input_line = read_line () in
      if input_line = "exit" then
        print_endline "再见!"
      else
        try
          (* 求值并打印结果 *)
          let result = s_expr_to_value input_line env in
          print_endline (string_of_value result);
          loop ()  (* 继续循环 *)
        with e ->
          Printf.printf "求值错误: %s\n" (Printexc.to_string e);
          loop ()
    with End_of_file ->  (* 处理 Ctrl+D *)
      print_endline "\n再见!";
      exit 0
  in
  loop ()  (* 启动REPL循环 *)

(* 主函数 *)
let () =
  ignore (test_builtins ());  (* 运行测试套件 *)
  repl()
